import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from keras.models import Sequential
from keras.layers import LSTM, Dense
import matplotlib.pyplot as plt

# 读取数据
data = pd.read_csv("your_data.csv")

# 选择含沙量作为目标变量
target_variable = '含沙量'
target_data = data[target_variable].values.reshape(-1, 1)

# 归一化数据
scaler = MinMaxScaler(feature_range=(0, 1))
target_data_normalized = scaler.fit_transform(target_data)

# 将数据转换为时间序列数据
def create_time_series_dataset(dataset, time_steps=1):
    X, y = [], []
    for i in range(len(dataset) - time_steps):
        a = dataset[i:(i + time_steps), 0]
        X.append(a)
        y.append(dataset[i + time_steps, 0])
    return np.array(X), np.array(y)

# 设置时间步长（根据需要调整）
time_steps = 3

# 创建时间序列数据集
X, y = create_time_series_dataset(target_data_normalized, time_steps)

# 将数据集分为训练集和测试集
train_size = int(len(X) * 0.8)
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]

# 将数据整形为LSTM模型所需的三维形状 (samples, time steps, features)
X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))

# 构建LSTM模型
model = Sequential()
model.add(LSTM(units=50, return_sequences=True, input_shape=(X_train.shape[1], 1)))
model.add(LSTM(units=50))
model.add(Dense(units=1))
model.compile(optimizer='adam', loss='mean_squared_error')

# 训练模型
model.fit(X_train, y_train, epochs=50, batch_size=32)

# 在测试集上进行预测
predictions = model.predict(X_test)
predictions = scaler.inverse_transform(predictions)

# 反归一化测试集目标变量
y_test = scaler.inverse_transform(y_test.reshape(-1, 1))

# 评估模型
rmse = np.sqrt(mean_squared_error(y_test, predictions))
print(f'RMSE: {rmse}')

# 绘制预测结果
plt.plot(y_test, label='Actual')
plt.plot(predictions, label='Predicted')
plt.legend()
plt.show()
